# 多项式升维(线性回归)
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import mean_squared_error

# 导 入 mean_squared_error 的 目 的 是 为 了 去 评 估 一 下 升 维的效果，
# 导入主角PolynomialFeatures 来做多项式回归升维，当然可以看出来它也没叫回归，而是叫什么Features，
# 很明显区别于那些回归算法，因为目的不同嘛，回归算法目的是为了去拟合，升维只是对数据集 X 进行变换。

np.random.seed(42)

m = 100
X = 6 * np.random.rand(m, 1) - 3
y = 0.5 * X ** 2 + X + 2 + np.random.rand(m, 1)

# 拆分成测试集和训练集
X_train = X[:80]
X_test = X[80:]
y_train = y[:80]
y_test = y[80:]

# plt.plot(X, y, 'b.')
# plt.show()

# 设置一个字典，用来保存维度与该维度(degree)对应的图像参数（颜色和性质）
dict_ = {1: 'g-', 2: 'r+', 10: 'y*'}
for i in dict_:
    # degree 为实际 i + 1
    f = PolynomialFeatures(degree=i, include_bias=True)
    X_train_f = f.fit_transform(X_train)
    X_test_f = f.transform(X_test)
    # print(X_train[0])
    # print(X_train_f[0])
    # print(X_train.shape)
    # print(X_train_f.shape)
    # 线性回归 PolynomialFeatures 中的 include_bias 设置为True 的话，LinearRegression 中的 fit_intercept 就不需要设置为 True 了，反之亦然
    lin = LinearRegression(fit_intercept=False)
    lin.fit(X_train_f, y_train)
    print('截距：%s , 系数：%s' % (lin.intercept_, lin.coef_))
    y_train_predict = lin.predict(X_train_f)
    y_test_predict = lin.predict(X_test_f)
    plt.plot(X_train_f[:, 1], y_train_predict, dict_[i])
    # 评估结果
    train_mse = mean_squared_error(y_train, y_train_predict)
    test_mse = mean_squared_error(y_test, lin.predict(X_test_f))
    print(f"Degree {i}: Train MSE={train_mse:.2f}, Test MSE={test_mse:.2f}")
plt.show()
